import os

import yaml
from langchain_openai import AzureChatOpenAI

from fin_embedding import MilvusStore
from fin_preprocess import DocPreprocessing
from fin_rag import RAG

with open("config.yaml", "r") as f:
    config = yaml.safe_load(f)


# Step1: 文件预处理, 默认带OCR
file_path = "data/new"
file_list = [os.path.join(file_path, x) for x in os.listdir(file_path)]
docs = DocPreprocessing().batch_process(file_list)

# Step2: 根据文件构建向量数据库

milvus_store = MilvusStore(
    embedding_model_name_or_path=config.get("embedding_model_name_or_path"),
    milvus_conn_args=config.get("milvus_conn_args"),
)
collection_name = "fin_emb"

# # 第一次构建的时候需要
# retriever = milvus_store.advance_build(
#     collection_name,
#     docs,
# )

# 构建完成后可以直接加载
retriever = milvus_store.get_retriever(collection_name)
lc_az_gpt35 = AzureChatOpenAI(**config.get("gpt35"))

# Step3: 进行检索和生成
print(RAG(retriever, lc_az_gpt35).get_retrival_result("基金的合同生效日是什么时候?"))
